#include "KDTreeTensor.h"

//******************************************************************
//FUNCTION:
torch::Tensor CPointsTensor::scales() {
    // Compute scales by finding the average of the three nearest neighbors for each point
    const auto Index = getIndex<KdTreeTensor>();
    torch::Tensor Scales = torch::zeros({ static_cast<long int>(m_Tensor.size(0)), 1 }, torch::kFloat32);
    const int Count = 4;

    std::vector<size_t> Indices(Count);
    std::vector<float> SqrDists(Count);
    for (size_t i = 0; i < m_Tensor.size(0); i++) {
        Index->knnSearch(reinterpret_cast<float*>(m_Tensor[i].data_ptr()), Count, Indices.data(), SqrDists.data());

        float Sum = 0.0;
        for (size_t j = 1; j < Count; j++) {
            Sum += std::sqrt(SqrDists[j]);
        }
        Scales[i] = Sum / (Count - 1);
    }

    return Scales;
}

CPointsTensor::~CPointsTensor() {
    freeIndex<KdTreeTensor>();
}